# examples/quantitative_strategies.py
"""
量化交易策略示例
实现多种经典的量化交易算法
"""

import sys
import numpy as np
import pandas as pd
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional
import warnings
warnings.filterwarnings('ignore')

# 添加项目路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.services.market_data_service import MarketDataService

class TradingStrategy:
    """交易策略基类"""
    
    def __init__(self, name: str):
        self.name = name
        self.signals = pd.DataFrame()
        self.returns = pd.DataFrame()
        self.positions = pd.DataFrame()
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成交易信号 - 需要在子类中实现"""
        raise NotImplementedError
        
    def backtest(self, data: pd.DataFrame, initial_capital: float = 100000) -> Dict:
        """回测策略"""
        # 生成信号
        signals = self.generate_signals(data)
        
        # 计算持仓
        positions = signals.shift(1)  # 延迟一天执行
        positions = positions.fillna(0)
        
        # 计算收益
        returns = data['close'].pct_change()
        strategy_returns = positions * returns
        
        # 累计收益
        cumulative_returns = (1 + strategy_returns).cumprod()
        
        # 计算绩效指标
        total_return = cumulative_returns.iloc[-1] - 1
        annual_return = (1 + strategy_returns.mean()) ** 252 - 1
        volatility = strategy_returns.std() * np.sqrt(252)
        sharpe_ratio = annual_return / volatility if volatility > 0 else 0
        
        # 最大回撤
        max_drawdown = (cumulative_returns / cumulative_returns.expanding().max() - 1).min()
        
        # 胜率
        winning_trades = strategy_returns[strategy_returns > 0]
        win_rate = len(winning_trades) / len(strategy_returns[strategy_returns != 0]) if len(strategy_returns[strategy_returns != 0]) > 0 else 0
        
        return {
            'total_return': total_return,
            'annual_return': annual_return,
            'volatility': volatility,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'win_rate': win_rate,
            'cumulative_returns': cumulative_returns,
            'strategy_returns': strategy_returns,
            'signals': signals,
            'positions': positions
        }

class MovingAverageCrossover(TradingStrategy):
    """移动平均交叉策略"""
    
    def __init__(self, short_window: int = 20, long_window: int = 50):
        super().__init__(f"MA_Crossover_{short_window}_{long_window}")
        self.short_window = short_window
        self.long_window = long_window
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成移动平均交叉信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算移动平均
        signals['short_ma'] = data['close'].rolling(self.short_window).mean()
        signals['long_ma'] = data['close'].rolling(self.long_window).mean()
        
        # 生成信号
        signals['signal'][self.long_window:] = np.where(
            signals['short_ma'][self.long_window:] > signals['long_ma'][self.long_window:], 1.0, 0.0
        )
        
        # 计算信号变化
        signals['positions'] = signals['signal'].diff()
        
        return signals['signal']

class RSIMeanReversion(TradingStrategy):
    """RSI均值回归策略"""
    
    def __init__(self, rsi_period: int = 14, oversold: float = 30, overbought: float = 70):
        super().__init__(f"RSI_MeanReversion_{rsi_period}_{oversold}_{overbought}")
        self.rsi_period = rsi_period
        self.oversold = oversold
        self.overbought = overbought
        
    def calculate_rsi(self, prices: pd.Series) -> pd.Series:
        """计算RSI指标"""
        delta = prices.diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        avg_gain = gain.rolling(self.rsi_period).mean()
        avg_loss = loss.rolling(self.rsi_period).mean()
        
        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
        
        return rsi
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成RSI均值回归信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算RSI
        signals['rsi'] = self.calculate_rsi(data['close'])
        
        # 生成信号
        signals.loc[signals['rsi'] < self.oversold, 'signal'] = 1.0  # 买入信号
        signals.loc[signals['rsi'] > self.overbought, 'signal'] = -1.0  # 卖出信号
        
        return signals['signal']

class BollingerBandsMeanReversion(TradingStrategy):
    """布林带均值回归策略"""
    
    def __init__(self, period: int = 20, std_dev: float = 2):
        super().__init__(f"BollingerBands_{period}_{std_dev}")
        self.period = period
        self.std_dev = std_dev
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成布林带均值回归信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算布林带
        signals['ma'] = data['close'].rolling(self.period).mean()
        signals['std'] = data['close'].rolling(self.period).std()
        signals['upper_band'] = signals['ma'] + (signals['std'] * self.std_dev)
        signals['lower_band'] = signals['ma'] - (signals['std'] * self.std_dev)
        
        # 生成信号
        signals.loc[data['close'] < signals['lower_band'], 'signal'] = 1.0  # 价格触及下轨，买入
        signals.loc[data['close'] > signals['upper_band'], 'signal'] = -1.0  # 价格触及上轨，卖出
        signals.loc[data['close'] == signals['ma'], 'signal'] = 0.0  # 价格回到中轨，平仓
        
        return signals['signal']

class MomentumStrategy(TradingStrategy):
    """动量策略"""
    
    def __init__(self, lookback_period: int = 20, holding_period: int = 5):
        super().__init__(f"Momentum_{lookback_period}_{holding_period}")
        self.lookback_period = lookback_period
        self.holding_period = holding_period
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成动量信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算动量
        signals['momentum'] = data['close'].pct_change(self.lookback_period)
        
        # 生成信号
        signals.loc[signals['momentum'] > 0, 'signal'] = 1.0  # 正动量，买入
        signals.loc[signals['momentum'] <= 0, 'signal'] = -1.0  # 负动量，卖出
        
        return signals['signal']

class MACDStrategy(TradingStrategy):
    """MACD策略"""
    
    def __init__(self, fast_period: int = 12, slow_period: int = 26, signal_period: int = 9):
        super().__init__(f"MACD_{fast_period}_{slow_period}_{signal_period}")
        self.fast_period = fast_period
        self.slow_period = slow_period
        self.signal_period = signal_period
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成MACD信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算MACD
        ema_fast = data['close'].ewm(span=self.fast_period).mean()
        ema_slow = data['close'].ewm(span=self.slow_period).mean()
        signals['macd'] = ema_fast - ema_slow
        signals['macd_signal'] = signals['macd'].ewm(span=self.signal_period).mean()
        signals['macd_histogram'] = signals['macd'] - signals['macd_signal']
        
        # 生成交易信号
        signals['signal'][1:] = np.where(
            (signals['macd'][1:] > signals['macd_signal'][1:]) & 
            (signals['macd'][:-1].values <= signals['macd_signal'][:-1].values), 1.0, 0.0
        )
        
        signals['signal'][1:] = np.where(
            (signals['macd'][1:] < signals['macd_signal'][1:]) & 
            (signals['macd'][:-1].values >= signals['macd_signal'][:-1].values), -1.0, 
            signals['signal'][1:]
        )
        return signals['signal']

class PairsTradingStrategy(TradingStrategy):
    """配对交易策略"""
    
    def __init__(self, stock1: str, stock2: str, lookback_period: int = 60, entry_threshold: float = 2.0, exit_threshold: float = 0.5):
        super().__init__(f"PairsTrading_{stock1}_{stock2}")
        self.stock1 = stock1
        self.stock2 = stock2
        self.lookback_period = lookback_period
        self.entry_threshold = entry_threshold
        self.exit_threshold = exit_threshold
        
    def generate_signals(self, data1: pd.DataFrame, data2: pd.DataFrame) -> pd.DataFrame:
        """生成配对交易信号"""
        # 确保数据对齐
        combined_data = pd.concat([data1['close'], data2['close']], axis=1, keys=[self.stock1, self.stock2])
        combined_data = combined_data.dropna()
        
        signals = pd.DataFrame(index=combined_data.index)
        signals['signal_stock1'] = 0.0
        signals['signal_stock2'] = 0.0
        
        # 计算价格比率
        signals['ratio'] = combined_data[self.stock1] / combined_data[self.stock2]
        
        # 计算移动平均和标准差
        signals['ratio_ma'] = signals['ratio'].rolling(self.lookback_period).mean()
        signals['ratio_std'] = signals['ratio'].rolling(self.lookback_period).std()
        
        # 计算Z-score
        signals['z_score'] = (signals['ratio'] - signals['ratio_ma']) / signals['ratio_std']
        
        # 生成交易信号
        # 当Z-score > entry_threshold时，做空股票1，做多股票2
        long_entry = signals['z_score'] > self.entry_threshold
        signals.loc[long_entry, 'signal_stock1'] = -1.0
        signals.loc[long_entry, 'signal_stock2'] = 1.0
        
        # 当Z-score < -entry_threshold时，做多股票1，做空股票2
        short_entry = signals['z_score'] < -self.entry_threshold
        signals.loc[short_entry, 'signal_stock1'] = 1.0
        signals.loc[short_entry, 'signal_stock2'] = -1.0
        
        # 当Z-score回归到exit_threshold附近时平仓
        exit_condition = np.abs(signals['z_score']) < self.exit_threshold
        signals.loc[exit_condition, 'signal_stock1'] = 0.0
        signals.loc[exit_condition, 'signal_stock2'] = 0.0
        
        return signals[['signal_stock1', 'signal_stock2']]

class MeanReversionStrategy(TradingStrategy):
    """均值回归策略（基于统计套利）"""
    
    def __init__(self, lookback_period: int = 20, entry_threshold: float = 2.0, exit_threshold: float = 0.5):
        super().__init__(f"MeanReversion_{lookback_period}_{entry_threshold}")
        self.lookback_period = lookback_period
        self.entry_threshold = entry_threshold
        self.exit_threshold = exit_threshold
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成均值回归信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算移动平均和标准差
        signals['ma'] = data['close'].rolling(self.lookback_period).mean()
        signals['std'] = data['close'].rolling(self.lookback_period).std()
        
        # 计算Z-score
        signals['z_score'] = (data['close'] - signals['ma']) / signals['std']
        
        # 生成信号
        # 价格偏离均值过多时进行反向操作
        signals.loc[signals['z_score'] > self.entry_threshold, 'signal'] = -1.0  # 卖出
        signals.loc[signals['z_score'] < -self.entry_threshold, 'signal'] = 1.0   # 买入
        signals.loc[np.abs(signals['z_score']) < self.exit_threshold, 'signal'] = 0.0  # 平仓
        
        return signals['signal']

class BreakoutStrategy(TradingStrategy):
    """突破策略"""
    
    def __init__(self, lookback_period: int = 20, breakout_threshold: float = 0.02):
        super().__init__(f"Breakout_{lookback_period}_{breakout_threshold}")
        self.lookback_period = lookback_period
        self.breakout_threshold = breakout_threshold
        
    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """生成突破信号"""
        signals = pd.DataFrame(index=data.index)
        signals['signal'] = 0.0
        
        # 计算阻力位和支撑位
        signals['resistance'] = data['high'].rolling(self.lookback_period).max()
        signals['support'] = data['low'].rolling(self.lookback_period).min()
        
        # 计算突破条件
        upward_breakout = (data['close'] > signals['resistance'] * (1 + self.breakout_threshold))
        downward_breakout = (data['close'] < signals['support'] * (1 - self.breakout_threshold))
        
        # 生成信号
        signals.loc[upward_breakout, 'signal'] = 1.0   # 向上突破，买入
        signals.loc[downward_breakout, 'signal'] = -1.0  # 向下突破，卖出
        
        return signals['signal']

class PortfolioOptimization:
    """投资组合优化"""
    
    @staticmethod
    def calculate_portfolio_metrics(returns: pd.DataFrame, weights: np.ndarray) -> Dict:
        """计算投资组合指标"""
        portfolio_return = np.sum(returns.mean() * weights) * 252
        portfolio_std = np.sqrt(np.dot(weights.T, np.dot(returns.cov() * 252, weights)))
        sharpe_ratio = portfolio_return / portfolio_std if portfolio_std > 0 else 0
        
        return {
            'return': portfolio_return,
            'volatility': portfolio_std,
            'sharpe_ratio': sharpe_ratio
        }
    
    @staticmethod
    def optimize_portfolio(returns: pd.DataFrame, target_return: Optional[float] = None) -> Dict:
        """投资组合优化（最大夏普比率）"""
        from scipy.optimize import minimize
        
        n_assets = len(returns.columns)
        
        # 目标函数：最小化负夏普比率
        def objective(weights):
            portfolio_return = np.sum(returns.mean() * weights) * 252
            portfolio_std = np.sqrt(np.dot(weights.T, np.dot(returns.cov() * 252, weights)))
            return -(portfolio_return / portfolio_std) if portfolio_std > 0 else -999
        
        # 约束条件
        constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})  # 权重和为1
        bounds = tuple((0, 1) for _ in range(n_assets))  # 权重在0-1之间
        
        # 初始权重
        initial_guess = np.array([1/n_assets] * n_assets)
        
        # 优化
        result = minimize(objective, initial_guess, method='SLSQP', 
                         bounds=bounds, constraints=constraints)
        
        if result.success:
            optimal_weights = result.x
            metrics = PortfolioOptimization.calculate_portfolio_metrics(returns, optimal_weights)
            
            return {
                'weights': optimal_weights,
                'expected_return': metrics['return'],
                'volatility': metrics['volatility'],
                'sharpe_ratio': metrics['sharpe_ratio'],
                'success': True
            }
        else:
            return {'success': False, 'message': 'Optimization failed'}

class StrategyComparer:
    """策略比较器"""
    
    def __init__(self):
        self.results = {}
        
    def add_strategy_result(self, strategy_name: str, result: Dict):
        """添加策略结果"""
        self.results[strategy_name] = result
        
    def compare_strategies(self) -> pd.DataFrame:
        """比较策略表现"""
        comparison_data = []
        
        for strategy_name, result in self.results.items():
            comparison_data.append({
                'Strategy': strategy_name,
                'Total Return': f"{result['total_return']:.2%}",
                'Annual Return': f"{result['annual_return']:.2%}",
                'Volatility': f"{result['volatility']:.2%}",
                'Sharpe Ratio': f"{result['sharpe_ratio']:.3f}",
                'Max Drawdown': f"{result['max_drawdown']:.2%}",
                'Win Rate': f"{result['win_rate']:.2%}"
            })
        
        return pd.DataFrame(comparison_data)
    
    def plot_cumulative_returns(self):
        """绘制累计收益对比图"""
        try:
            import matplotlib.pyplot as plt
            
            plt.figure(figsize=(12, 8))
            
            for strategy_name, result in self.results.items():
                if 'cumulative_returns' in result:
                    plt.plot(result['cumulative_returns'].index, 
                            result['cumulative_returns'], 
                            label=strategy_name)
            
            plt.title('策略累计收益对比')
            plt.xlabel('日期')
            plt.ylabel('累计收益')
            plt.legend()
            plt.grid(True)
            plt.show()
            
        except ImportError:
            print("⚠️ matplotlib未安装，无法绘制图表")

def main():
    """主函数演示量化交易策略"""
    
    print("📊 量化交易策略回测系统")
    print("=" * 50)
    
    # 初始化服务
    market_service = MarketDataService()
    comparer = StrategyComparer()
    
    # 获取测试数据
    symbol = 'AAPL'
    print(f"📈 获取 {symbol} 历史数据...")
    
    start_date = datetime.now() - timedelta(days=365*2)  # 2年数据
    end_date = datetime.now()
    
    data = market_service.get_stock_price(
        symbol,
        start_date=start_date,
        end_date=end_date,
        interval='1d'
    )
    
    if data.empty:
        print("❌ 未能获取数据")
        return
    
    print(f"✅ 获取到 {len(data)} 条数据记录")
    
    # 定义策略列表
    strategies = [
        MovingAverageCrossover(20, 50),
        MovingAverageCrossover(10, 30),
        RSIMeanReversion(14, 30, 70),
        BollingerBandsMeanReversion(20, 2),
        MomentumStrategy(20, 5),
        MACDStrategy(12, 26, 9),
        MeanReversionStrategy(20, 2.0, 0.5),
        BreakoutStrategy(20, 0.02)
    ]
    
    print(f"\n🔬 开始回测 {len(strategies)} 个策略...")
    
    # 回测每个策略
    for strategy in strategies:
        try:
            print(f"   回测策略: {strategy.name}")
            result = strategy.backtest(data)
            comparer.add_strategy_result(strategy.name, result)
            
            print(f"     ✅ 总收益: {result['total_return']:.2%}, "
                  f"夏普比率: {result['sharpe_ratio']:.3f}")
            
        except Exception as e:
            print(f"     ❌ 策略 {strategy.name} 回测失败: {e}")
    
    # 买入持有策略作为基准
    print(f"   回测基准策略: Buy & Hold")
    buy_hold_returns = data['close'].pct_change().fillna(0)
    buy_hold_cumulative = (1 + buy_hold_returns).cumprod()
    
    buy_hold_result = {
        'total_return': buy_hold_cumulative.iloc[-1] - 1,
        'annual_return': (1 + buy_hold_returns.mean()) ** 252 - 1,
        'volatility': buy_hold_returns.std() * np.sqrt(252),
        'sharpe_ratio': ((1 + buy_hold_returns.mean()) ** 252 - 1) / (buy_hold_returns.std() * np.sqrt(252)),
        'max_drawdown': (buy_hold_cumulative / buy_hold_cumulative.expanding().max() - 1).min(),
        'win_rate': len(buy_hold_returns[buy_hold_returns > 0]) / len(buy_hold_returns[buy_hold_returns != 0]),
        'cumulative_returns': buy_hold_cumulative
    }
    
    comparer.add_strategy_result('Buy_Hold', buy_hold_result)
    
    # 显示策略比较结果
    print(f"\n📊 策略表现对比:")
    comparison_df = comparer.compare_strategies()
    print(comparison_df.to_string(index=False))
    
    # 投资组合优化示例
    print(f"\n🎯 投资组合优化示例...")
    
    # 获取多只股票数据
    symbols = ['AAPL', 'GOOGL', 'MSFT', 'TSLA']
    print(f"获取投资组合股票数据: {', '.join(symbols)}")
    
    portfolio_data = market_service.get_stock_price(
        symbols,
        start_date=start_date,
        end_date=end_date,
        interval='1d'
    )
    
    if len(portfolio_data) >= 2:
        # 构建收益率矩阵
        returns_data = {}
        for symbol, stock_data in portfolio_data.items():
            if not stock_data.empty:
                returns_data[symbol] = stock_data['close'].pct_change()
        
        returns_df = pd.DataFrame(returns_data).dropna()
        
        if not returns_df.empty:
            # 投资组合优化
            optimization_result = PortfolioOptimization.optimize_portfolio(returns_df)
            
            if optimization_result['success']:
                print(f"✅ 投资组合优化成功:")
                print(f"   预期年化收益: {optimization_result['expected_return']:.2%}")
                print(f"   预期年化波动: {optimization_result['volatility']:.2%}")
                print(f"   夏普比率: {optimization_result['sharpe_ratio']:.3f}")
                print(f"   最优权重:")
                
                for i, symbol in enumerate(returns_df.columns):
                    weight = optimization_result['weights'][i]
                    print(f"     {symbol}: {weight:.1%}")
            else:
                print("❌ 投资组合优化失败")
    
    # 配对交易示例
    if len(portfolio_data) >= 2:
        print(f"\n🔄 配对交易策略示例...")
        symbols_list = list(portfolio_data.keys())
        stock1, stock2 = symbols_list[0], symbols_list[1]
        
        if not portfolio_data[stock1].empty and not portfolio_data[stock2].empty:
            pairs_strategy = PairsTradingStrategy(stock1, stock2, lookback_period=60)
            
            try:
                pairs_signals = pairs_strategy.generate_signals(
                    portfolio_data[stock1], 
                    portfolio_data[stock2]
                )
                
                print(f"✅ {stock1}-{stock2} 配对交易信号生成成功")
                print(f"   信号数量: {len(pairs_signals[pairs_signals != 0])}")
                
                # 显示最近的交易信号
                recent_signals = pairs_signals.tail(10)
                non_zero_signals = recent_signals[(recent_signals != 0).any(axis=1)]
                if not non_zero_signals.empty:
                    print(f"   最近的交易信号:")
                    print(non_zero_signals)
                    
            except Exception as e:
                print(f"❌ 配对交易策略失败: {e}")
    
    print(f"\n🎯 策略总结:")
    print(f"   测试了 {len(strategies)} 个单资产策略")
    print(f"   包含了技术分析、动量、均值回归等多种类型")
    print(f"   提供了投资组合优化和配对交易的高级策略")
    print(f"   可以根据实际需求调整参数和组合策略")
    
    # 绘制累计收益图（如果matplotlib可用）
    try:
        comparer.plot_cumulative_returns()
    except:
        print(f"\n💡 提示: 安装matplotlib可查看策略收益对比图")

if __name__ == "__main__":
    main()     
        
